load("nelson_data.RData")
attach(kkk)

T <- as.matrix(vidcdum)
M <- as.matrix(idisord)
Y <- as.matrix(kspeech)

err.cr <- cor(lm(M ~ T)$resid, lm(Y ~ T)$resid)

#Start Simulation Run
n <- length(Y)
rho <- round(seq(-.95, .95, by = .01),2)
med.eff <- matrix(NA, length(rho), 1)
med.var <- matrix(NA, length(rho), 4)
b2 <- matrix(NA, length(rho), 1)

for(i in 1:length(rho)){

e.cor <- rho[i]

b.dif <- 1
eps <- .Machine$double.eps

#Stacked Equations
Y.c <- rbind(M,Y)
cons <- matrix(1, n, 1)
X.1 <- cbind(cons, T, matrix(0,n, 3))
X.2 <- cbind(matrix(0,n,2), cons, T, M)
X <- rbind(X.1, X.2)

#Estimates of OLS Start Values
inxx <- solve(crossprod(X))
b.ols <- inxx %*% crossprod(X,Y.c)
b.tmp <- b.ols

while(abs(b.dif) > eps){

e.hat <- as.matrix(Y.c - (X %*% b.tmp))

e.1 <- e.hat[1:n]
e.2 <- e.hat[(n+1): (2*n)]

sd.1 <- sd(e.1)
sd.2 <- sd(e.2)

omega <- matrix(NA, 2,2)

omega[1,1] <- crossprod(e.1)/(n-2)
omega[2,2] <- crossprod(e.2)/(n-3) 
omega[2,1] <- e.cor*sd.1*sd.2
omega[1,2] <- e.cor*sd.1*sd.2

I <- diag(1,n)
omega.i <- solve(omega)
v.i <-  kronecker(omega.i, I)

X.sur <- crossprod(X, v.i) %*% X 
b.sur <- solve(X.sur) %*% crossprod(X, v.i) %*% Y.c

#Variance-Covariance Matrix
v.cov <- crossprod(X, v.i) %*% X
v.cov <- solve(v.cov)

b.old <- b.tmp
b.dif <- sum((b.sur - b.old)^2)
b.tmp <- b.sur

}


med.eff[i,] <- b.sur[2]*b.sur[5]
b2[i,] <- b.sur[2]

#BK
med.var[i,1] <- (b.sur[2]^2 * v.cov[5,5]) + (b.sur[5]^2*v.cov[2,2]) + v.cov[5,5]*v.cov[2,2]
#Sobel
med.var[i,2] <- (b.sur[2]^2 * v.cov[5,5]) + (b.sur[5]^2*v.cov[2,2])
#Product of Normals 
med.var[i,3] <- (b.sur[2]^2 * v.cov[5,5]) + (b.sur[5]^2*v.cov[2,2]) + (2*b.sur[2]*b.sur[5]*v.cov[5,2]) + v.cov[2,2]*v.cov[5,5] + v.cov[5,2]^2
#Delta Method
med.var[i,4] <-  (b.sur[2]^2 * v.cov[5,5]) + (b.sur[5]^2*v.cov[2,2]) + (2*b.sur[2]*b.sur[5]*v.cov[5,2])

}

upper <- med.eff + qnorm(0.975) * sqrt(med.var[,4])
lower <- med.eff - qnorm(0.975) * sqrt(med.var[,4])

pdf("sp-iordII.pdf", width=4, height=4, onefile=FALSE, paper="special")
plot(rho, med.eff, type="n", xlab="", ylab = "", ylim = c(-3,1.25))
polygon(x=c(rho, rev(rho)), y=c(lower, rev(upper)), border=FALSE, col=8, lty=2)
lines(rho, med.eff, lty=1)
abline(h=0)
abline(h=med.eff[96], lty=2)
abline(v=0)
title(xlab=expression(paste("Sensitivity Parameter: ", rho)), line=2.5, cex.lab=.9)
title(ylab = expression(paste("Average Mediation Effect: ", bar(delta))), cex.lab=.9)
dev.off()

